import type { Key, TreeLayout, TreeLike } from "@/models";
import { Position, type Edge, type Elements, type Node } from "@vue-flow/core";
import { cloneDeep } from "lodash-es";

export const DEFAULT_LAYOUT: TreeLayout = { mode: "vertical", gap: { x: 10, y: 50 }, size: { w: 150, h: 36 } };

export class CalcTree {
  flat: TreeLike[] = [];
  trees: TreeLike[] = [];
  pools: Map<Key, TreeLike[][]> = new Map();

  constructor(flat: TreeLike[]) {
    this.flat = flat;
    this.toTrees();
    this.toPools();
  }

  private toTrees() {
    this.flat.forEach((tl) => {
      const parent = this.flat.find((p) => p.id === tl.pid);
      if (parent) {
        if (!parent.children) parent.children = [];
        parent.children.push(tl);
        tl.parent = parent;
      } else {
        this.trees.push(tl);
      }
    });
  }

  private toPools() {
    function toPool(tree: TreeLike) {
      const pool: TreeLike[][] = [];

      function send2pool(tl: TreeLike, depth: number) {
        tl.depth = depth;
        if (pool[depth] === undefined) pool[depth] = [];
        pool[depth].push(tl);

        if (tl.children) {
          // 真子集
          for (const child of tl.children) {
            send2pool(child, depth + 1);
          }
        } else {
          // 虚拟子集，占位用
          if (hasNephew(tl)) {
            const virtualChild: TreeLike = { id: "", name: "", pid: tl.id, parent: tl, virtual: true };
            tl.children = [virtualChild];
            send2pool(virtualChild, depth + 1);
          }
        }
      }

      // 是否有侄辈
      function hasNephew(tl: TreeLike) {
        const brothers = findBrothers(tl);
        let found = false;

        for (const brother of brothers) {
          if (brother.children) {
            if (brother.children.findIndex((c) => !c.virtual) > -1) {
              found = true;
              break;
            }
          }
        }

        return found;
      }

      // 查找同辈，含远亲同辈
      function findBrothers(tl: TreeLike) {
        const brothers: TreeLike[] = [];

        if (tl.parent) {
          const p = tl.parent;
          brothers.push(...p.children!.filter((c) => c.id !== tl.id));

          const uncles = findBrothers(tl.parent);
          for (const uncle of uncles) {
            if (uncle.children) brothers.push(...uncle.children);
          }
        }

        return brothers;
      }

      send2pool(tree, 0);

      // set sequence
      const max = pool[pool.length - 1].length;
      for (let i = pool.length - 1; i >= 0; i--) {
        const row = pool[i];
        const lns: number[] = [];
        for (let j = 0; j < row.length; j++) {
          const cell = row[j];
          if (!cell.children) {
            cell.length = 1 / max;
          } else {
            cell.length = cell.children.map((c) => c.length || 1).reduce((a, b) => a + b);
          }
          lns.push(cell.length);

          const prefix = j > 0 ? lns.slice(0, j).reduce((a, b) => a + b) : 0;
          cell.sequence = prefix + cell.length / 2;
        }
      }

      return pool;
    }

    for (const tree of this.trees) {
      this.pools.set(tree.id, toPool(tree));
    }
  }
}

export function generateElements(flat: TreeLike[], layout?: Partial<TreeLayout>) {
  const _layout: TreeLayout = Object.assign(DEFAULT_LAYOUT, layout);
  const _flat = cloneDeep(flat);

  const elements: Elements = [];
  const pool: TreeLike[][] = new CalcTree(_flat).pools.values().next().value || [];
  const xmax = pool[pool.length - 1].length;
  const smax = _layout.mode === "vertical" ? xmax * _layout.size.w + (xmax - 1) * _layout.gap.x : xmax * _layout.size.h + (xmax - 1) * _layout.gap.y;

  for (let i = 0; i < pool.length; i++) {
    const row = pool[i];
    for (let j = 0; j < row.length; j++) {
      const cell = row[j];
      if (cell.virtual) continue;

      let node: Node<TreeLike>;

      if (_layout.mode === "vertical") {
        node = {
          id: cell.id + "",
          label: cell.name,
          sourcePosition: Position.Bottom,
          targetPosition: Position.Top,
          position: {
            x: smax * (cell.sequence || 0),
            y: (_layout.size.h + _layout.gap.y) * (cell.depth || 0),
          },
        };
      } else {
        node = {
          id: cell.id + "",
          label: cell.name,
          sourcePosition: Position.Right,
          targetPosition: Position.Left,
          position: {
            x: (_layout.size.w + _layout.gap.x) * (cell.depth || 0),
            y: smax * (cell.sequence || 0),
          },
        };
      }

      if (!cell.parent) node.type = "input";
      if (!cell.children || cell.children.every((c) => c.virtual)) node.type = "output";
      elements.push(node);

      if (cell.pid !== undefined) {
        const edge: Edge = {
          id: `e${cell.pid}-${cell.id}`,
          source: cell.pid + "",
          target: cell.id + "",
          type: "smoothstep",
          animated: true,
          style: { stroke: "orangered" },
        };
        elements.push(edge);
      }
    }
  }

  return elements;
}
